#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author  : jingcb
# @Time    : 2018/7/31 下午3:54
# @File    : read_cog.py
# @Project : project-5yuan

import os
from rasterio.transform import from_origin
from rasterio.windows import Window
import mercantile
import numpy
import rasterio
from rasterio.vrt import WarpedVRT
from rasterio.enums import Resampling
import json
import math
import subprocess
from PIL import Image

def read_by_zxy(filepath, z, x, y, src_crs='EPSG:3857'):
    nodata = 0
    mercator_tile = mercantile.Tile(x=x, y=y, z=z)  # line:157

    w, s, e, n = mercantile.bounds(mercator_tile)
    if src_crs == 'EPSG:3857':
        w, s = mercantile.xy(w, s)
        e, n = mercantile.xy(e, n)
    with rasterio.open(filepath) as src:
        indexes = [i for i in range(1, src.count+1)]
        with WarpedVRT(src,
                       dst_crs=src_crs,
                       resampling=Resampling.bilinear,
                       src_nodata=nodata,
                       dst_nodata=nodata) as vrt:
                            window = vrt.window(w, s, e, n, precision=21)
                            result = vrt.read(window=window,
                                            boundless=True,
                                            resampling=Resampling.nearest,
                                            out_shape=(256,256),
                                            indexes=indexes)
                            return result


def get_zoomres_3857(z):
    with open('./lod.json') as f:

        zoom_res = json.load(f)['zoomReses']
        return zoom_res[z]


def read_by_bbox(filepath,zoom_res,bbox,src_crs='EPSG:3857'):
    """

    :param filepath:
    :param z:
    :param bbox: list [w,s,e.n]
    :return:
    """
    nodata = 0
    w, s, e, n = bbox
    if src_crs == 'EPSG:3857':
        w, s = mercantile.xy(w, s)
        e, n = mercantile.xy(e, n)
    out_shape = (math.ceil(abs(n-s)/zoom_res),math.ceil(abs(w-e)/zoom_res))
    with rasterio.open(filepath) as src:
        indexes = [i for i in range(1, src.count + 1)]
        with WarpedVRT(src,
                       dst_crs=src_crs,
                       resampling=Resampling.bilinear,
                       src_nodata=nodata,
                       dst_nodata=nodata) as vrt:
            window = vrt.window(w, s, e, n, precision=21)
            result = vrt.read(window=window,
                              boundless=True,
                              resampling=Resampling.nearest,
                              out_shape=out_shape,
                              indexes=indexes)
            return result

def read_src_by_bbox(filepath,bbox):
    w, s, e, n = bbox
    with rasterio.open(filepath) as src :
        (start_col,start_row) = src.index(x=w,y=n)
        (end_col, end_row) = src.index(x=e, y=s)

        data = src.read(window=((start_row,end_row),(start_col,end_col)))
        return data


def _create_tif_with_src(src_path,dst_path, width,height,bands_count):
    with rasterio.open(src_path) as src:
        (lon,lat) = src.xy(0,0)
        lat_res = (src.bounds.top-src.bounds.bottom)/height
        lon_res = (src.bounds.right - src.bounds.left) /width
        transform = from_origin(lon - lon_res / 2, lat + lat_res / 2, lon_res, lat_res)
        dataset = rasterio.open(dst_path,'w',driver='GTiff',height=height,width=width,count=bands_count,dtype=rasterio.uint16,crs=src.crs,transform=transform)
        dataset.close()



def thumb(src_path,dst_path,mode="tif",colomap=None):
    """
    create thumb of src tif
    :param src_path: str;    input src tif path;
    :param dst_path: str;    output image path
    :param mode: str; tif or cog
    :param colomap: dict
    :return:
    """
    if mode == 'tif':
        addo_step = "gdaladdo -ro {} 2 4 8 16 32".format(src_path)
    subprocess.call(addo_step, shell=True)
    with rasterio.open(src_path) as src:

        with WarpedVRT(src, dst_crs="EPSG:4326", resampling=Resampling.bilinear,src_nodata = 0, dst_nodata = 0) as vrt:
            (w, s, e, n) = src.bounds
            bands_count = src.count
            height= 512

            width = int(src.width/(src.height/512))
            result = vrt.read(boundless=True,resampling=Resampling.nearest,out_shape=(height,width))



    if bands_count >=3:
        result = result[0:3]

        #拉伸到【1，255】 按照 2% 98%拉伸
        # for i in range(3):
        #     hist = numpy.percentile(result[i][result[i] > 0], (2, 98)).astype(numpy.int).tolist()
        #     result[i] = numpy.where(result[i] > 0,
        #                         _linear_rescale(result[i],
        #                                              in_range=[hist[0],hist[1]],
        #
        #                                              out_range=[1, 255]), 0)

    if len(result.shape) >= 3:
        result = numpy.ma.transpose(result, [1, 2, 0])
        result = result.squeeze()
    if bands_count == 1 and colomap != None:
        arr = numpy.zeros((height,width,3),dtype=numpy.uint8)

        for v_color in colomap.values():
            maskmin =  result>= v_color["value"][0]
            maskmax = result < v_color["value"][1]
            mask = maskmax & maskmin
            arr[numpy.where(mask == True)] = numpy.array(v_color["color"])

        result = arr
    if result.dtype != numpy.uint8:
        result = result.astype(numpy.uint8)
    img = Image.fromarray(result,mode="RGB")

    img.save(dst_path)





def write_cog(src_path, dst_path,width, height, data, bbox):
    """
    :param src_path:
    :param dst_path:
    :param width:
    :param height:
    :param data:
    :param bbox: [w,s,e,n]
    :return:
    """
    w, s, e, n = bbox
    (bands_count, data_h, data_w) = data.shape
    data = data.astype(numpy.uint16)
    if not os.path.exists(dst_path):
        _create_tif_with_src(src_path,dst_path,width,height,bands_count)
    with rasterio.open(dst_path,'r+') as dst:
        row_start, col_start = dst.index(w,n)
        row_end, col_end = dst.index(e,s)
        window = Window.from_slices((row_start, row_end), (col_start, col_end))
        dst.write(data,window=window)




if __name__ == '__main__':
    colormap = {
        "8": {
            "value": (1735, 1900),
            "color": [215, 25, 28],
        },
        "1":{
            "value": (39,163),
            "color": [43, 131,186]
        },
        "2": {
            "value": (163, 223),
            "color": [116, 183, 174]
        },
        "3": {
            "value": (223, 293),
            "color": [183, 226, 168],
        },
        "4": {
            "value": (293, 372),
            "color": [231, 246, 184],
        },
        "5": {
            "value": (372, 531),
            "color": [255, 232, 164],
        },
        "6": {
            "value": (531, 776),
            "color": [254, 186, 110],
        },
        "7": {
            "value": (776, 1735),
            "color": [237, 110, 67],
        }


    }
    thumb("/Users/jingcb/geohey/project-5yuan/data_tmp/test.tif","/Users/jingcb/geohey/project-5yuan/data_tmp/dem1.tif",colomap=colormap)
    #data = read_by_zxy('/Users/jingcb/geohey/project-5yuan/data_tmp/GF2_PMS1_E109.2_N19.8_20151027_L1A0001416673-MSS1.tiff.cog',6,51,28)
    # with rasterio.open('/Users/jingcb/geohey/project-5yuan/data_tmp/GF2_PMS1_E109.2_N19.8_20151027_L1A0001416673-MSS1.tiff.cog') as src:
    #     width, height = src.width/2, src.height/2
    #     bbox = [i for i in src.bounds]
    #     zoom_res = src.res
    # data = read_by_bbox('/Users/jingcb/geohey/project-5yuan/data_tmp/GF2_PMS1_E109.2_N19.8_20151027_L1A0001416673-MSS1.tiff.cog',zoom_res[0]*2,bbox,src_crs='EPSG:4326')
    # write_cog('/Users/jingcb/geohey/project-5yuan/data_tmp/GF2_PMS1_E109.2_N19.8_20151027_L1A0001416673-MSS1.tiff.cog','/Users/jingcb/geohey/project-5yuan/data_tmp/test.tif',width,height,data,bbox)